import os
import yaml
from types import SimpleNamespace
from collections import OrderedDict

from tqdm import tqdm
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torch.utils.data import DataLoader, WeightedRandomSampler, Dataset
from torch import optim
import torchvision.transforms as T
from sklearn import metrics
import matplotlib.pyplot as plt

from data.datasets import AortaDataset3DCenter, get_weight_list, AortaDataset
from utils.ranger import Ranger
from utils.lr_scheduler import CosineAnnealingWithWarmUpLR
from model.resnet3d import resnet3d
from model.SupCon import resnet
import data.transforms as MT

os.environ['CUDA_VISIBLE_DEVICES'] = '0'
device = torch.device('cuda:0')

#3d数据和网络
image_size = 81
source = ''
cate = ['0','1']
n_classes = len(cate)
sobel = True
vt_list = [
    MT.Resize3D(image_size),
    MT.CenterCrop3D(image_size),
    MT.ToTensor3D(),
]
if sobel:
    vt_list.append(MT.SobelChannel(3, flag_3d=True))
val_transform = T.Compose(vt_list)
test_dataset = AortaDataset3DCenter(os.path.join(source, 'test'), cate, val_transform, depth=7, step=1)
test_loader = DataLoader(test_dataset,
                         batch_size=128,
                         shuffle=False,
                         drop_last=False,
                         num_workers=8, 
                         pin_memory=True)

load_model = ''
net = resnet3d(34, n_channels=2 if sobel else 1, n_classes=n_classes, conv1_t_size=3)
net.load_state_dict(torch.load(load_model, map_location=device))
net.to(device)
net.eval()
print()

tot_loss = 0
true_list = []
pred_list = []
pred_ori_list = []
with tqdm(total=len(test_dataset), desc=f'test round', unit='img', leave=False) as pbar:
    for imgs, labels in test_loader:
        imgs, labels = imgs.to(device), labels.to(device)
        preds = net(imgs)
        if n_classes > 1:
            tot_loss += F.cross_entropy(preds, labels).item() * labels.size(0)
            pred_idx = torch.softmax(preds, dim=1)
        else:
            tot_loss += F.binary_cross_entropy_with_logits(preds.squeeze(1), labels.float()).item() * labels.size(0)
            pred_sig = torch.sigmoid(preds)
            pred_idx = torch.cat([1-pred_sig, pred_sig], dim=1)
        pred_ori_list += pred_idx.tolist()
        pred_idx = pred_idx.argmax(dim=1)
        labels_list = labels.tolist()
        true_list += labels_list
        pred_idx = pred_idx.tolist()
        pred_list.extend(pred_idx)
        pbar.update(labels.size(0))
tot_loss /= len(test_dataset)

print(f'Test report:\n'+metrics.classification_report(true_list, pred_list, digits=4))

metrics.roc_curve(true_list, pred_ori_list)
metrics.roc_auc_score(true_list, pred_ori_list)
